package test.demo;

import cn.hutool.core.util.ArrayUtil;
import com.alibaba.fastjson.JSON;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;
import tech.waterism.util.JacksonUtil;

import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 三水源滞后演算法（三水源汇流）
 * 系统模型集成
 *
 * @author 初京刚
 * @version 4.0
 */
public class LAG_3 extends PredictBase {

    private double C0;
    private double C1;
    private double C2;
    private double CI;
    private double CG;
    private double X;
    private double KK;
    private double LAG;
    private double CS;
    private int MP;
    private double F;
    private Date[] TM;
    private double[][] R;
    private int T;
    private double[] QSIG;
    private double[] QXSIG;
    private double[] QSIG0;
    private double QSP;
    private double QIP;
    private double QGP;
    /**
     * 预报开始时间
     */
    private Date startTm;

    /**
     * 预报开始数组下标
     */
    private int startIndex;

    public LAG_3(Map<String, String> data) {
        super();

        CG = Double.parseDouble(data.get("CG"));
        CI = Double.parseDouble(data.get("CI"));
        X = Double.parseDouble(data.get("X"));
        KK = Double.parseDouble(data.get("KK"));
        LAG = Double.parseDouble(data.get("LAG"));
        CS = Double.parseDouble(data.get("CS"));
        MP = Integer.parseInt(data.get("MP"));
        F = Double.parseDouble(data.get("F"));

        // 解析日期字符串数组
        Object dtObj = data.get("dt");
        String[] dtStrings;
        if (dtObj instanceof com.alibaba.fastjson.JSONArray) {
            com.alibaba.fastjson.JSONArray dtArray = (com.alibaba.fastjson.JSONArray) dtObj;
            dtStrings = new String[dtArray.size()];
            for (int i = 0; i < dtArray.size(); i++) {
                dtStrings[i] = dtArray.getString(i);
            }
        } else {
            dtStrings = JSON.parseObject(data.get("dt"), String[].class);
        }
        TM = new Date[dtStrings.length];
        for (int i = 0; i < dtStrings.length; i++) {
            try {
                // 使用SimpleDateFormat解析日期字符串
                java.text.SimpleDateFormat sdf = new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
                TM[i] = sdf.parse(dtStrings[i]);
            } catch (Exception e) {
                // 如果解析失败，使用当前时间
                TM[i] = new Date();
            }
        }
        
        // 解析开始时间
        Object startObj = data.get("start");
        String startString;
        if (startObj instanceof String) {
            startString = (String) startObj;
        } else if (startObj instanceof com.alibaba.fastjson.JSONArray) {
            com.alibaba.fastjson.JSONArray startArray = (com.alibaba.fastjson.JSONArray) startObj;
            startString = startArray.getString(0);
        } else {
            startString = JSON.parseObject(data.get("start"), String.class);
        }
        try {
            java.text.SimpleDateFormat sdf = new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
            startTm = sdf.parse(startString);
        } catch (Exception e) {
            // 如果解析失败，使用第一个时间
            startTm = TM[0];
        }
        startIndex = ArrayUtil.indexOf(TM, startTm);
        // 解析rSim数组
        Object rSimObj = data.get("rSim");
        if (rSimObj instanceof com.alibaba.fastjson.JSONArray) {
            com.alibaba.fastjson.JSONArray rSimArray = (com.alibaba.fastjson.JSONArray) rSimObj;
            R = new double[rSimArray.size()][];
            for (int i = 0; i < rSimArray.size(); i++) {
                com.alibaba.fastjson.JSONArray innerArray = rSimArray.getJSONArray(i);
                R[i] = new double[innerArray.size()];
                for (int j = 0; j < innerArray.size(); j++) {
                    R[i][j] = innerArray.getDoubleValue(j);
                }
            }
        } else {
            R = JSON.parseObject(data.get("rSim"), double[][].class);
        }
        
        T = Integer.parseInt(data.get("clen"));        // index.clen
        
        // 解析QSIG数组
        Object qsigObj = data.get("QSIG");
        if (qsigObj instanceof com.alibaba.fastjson.JSONArray) {
            com.alibaba.fastjson.JSONArray qsigArray = (com.alibaba.fastjson.JSONArray) qsigObj;
            QSIG0 = new double[qsigArray.size()];
            for (int i = 0; i < qsigArray.size(); i++) {
                QSIG0[i] = qsigArray.getDoubleValue(i);
            }
        } else {
            QSIG0 = JSON.parseObject(data.get("QSIG"), double[].class);
        }
        QSP = Double.parseDouble(data.get("QSP"));
        QIP = Double.parseDouble(data.get("QIP"));
        QGP = Double.parseDouble(data.get("QGP"));
        // 解析QXSIG数组
        Object qxsigObj = data.get("QXSIG");
        double[] qxsigArray;
        if (qxsigObj instanceof com.alibaba.fastjson.JSONArray) {
            com.alibaba.fastjson.JSONArray qxsigJsonArray = (com.alibaba.fastjson.JSONArray) qxsigObj;
            qxsigArray = new double[qxsigJsonArray.size()];
            for (int i = 0; i < qxsigJsonArray.size(); i++) {
                qxsigArray[i] = qxsigJsonArray.getDoubleValue(i);
            }
        } else {
            qxsigArray = JSON.parseObject(data.get("QXSIG"), double[].class);
        }
        
        QXSIG = new double[MP + 1];
        int I;
        for (I = 0; I < MP + 1; I++) {
            if (I > qxsigArray.length - 1) {
                QXSIG[I] = QXSIG[I - 1];
            } else {
                QXSIG[I] = qxsigArray[I];
            }
        }
    }

    /**
     * 模型计算
     */
    public Map<String, Object> LAG3model() {
        double[] RS = R[0];
        double[] RI = R[1];
        double[] RG = R[2];
        int Pnum = RS.length;
        QSIG = new double[Pnum + (int) LAG];

        if (LAG <= 1) {
            QSIG[0] = QSIG0[0];
        } else {
            int I;
            for (I = 0; I < LAG; I++) {
                if (QSIG0.length >= LAG) {
                    QSIG[I] = QSIG0[I];
                } else {
                    if (I < QSIG0.length) {
                        QSIG[I] = QSIG0[I];
                    } else {
                        QSIG[I] = QSIG[I - 1];
                    }
                }
            }
        }

        double CP = F / T / 3.6;
        CG = Math.pow(CG, T / 24.0);
        CI = Math.pow(CI, T / 24.0);
        double FKT = KK - KK * X + 0.5 * T;
        C0 = (0.5 * T - KK * X) / FKT;
        C1 = (KK * X + 0.5 * T) / FKT;
        C2 = (KK - KK * X - 0.5 * T) / FKT;

        double QSIG1;
        double[][] QSIGS = new double[Pnum + 1][QSIG0.length];
        if (LAG <= 1) {
            QSIG1 = QSIG[0];
        } else {
            QSIG1 = QSIG[(int) (LAG - 1)];
        }

        for (int i = 0; i < QSIG0.length; i++) {
            QSIGS[0][i] = QSIG0[i];
        }

        double[][] QXSIGS = new double[Pnum + 1][MP + 1];
        double[] QXSIGSS = new double[MP + 1];
        for (int I = 0; I < MP + 1; I++) {
            QXSIGS[0][I] = QXSIG[I];
            QXSIGSS[I] = QXSIG[I];
        }

        double[] QG = new double[Pnum];
        double[] QS = new double[Pnum];
        double[] QI = new double[Pnum];

        double[] QGS = new double[Pnum + 1];
        double[] QSS = new double[Pnum + 1];
        double[] QIS = new double[Pnum + 1];
        QGS[0] = QGP;
        QSS[0] = QSP;
        QIS[0] = QIP;


        // 陡降处理
        Calendar time0 = Calendar.getInstance();
        time0.setTime(TM[startIndex]);
        Integer Mth0 = time0.get(Calendar.MONTH) + 1; // 获得降雨开始月份
        Integer Yth0 = time0.get(Calendar.YEAR);

        for (int J = 0; J < Pnum; J++) {
            QGP = QGP * CG + RG[J] * (1.0 - CG) * CP;
            QIP = QIP * CI + RI[J] * (1.0 - CI) * CP;
            QSP = RS[J] * CP;

            QG[J] = QGP;
            QS[J] = QSP;
            QI[J] = QIP;

            QGS[J + 1] = QGP;
            QSS[J + 1] = QSP;
            QIS[J + 1] = QIP;

            QSIG1 = QSIG1 * CS + (QGP + QIP + QSP) * (1.0 - CS);
            double QTSIG = QSIG1;
            QSIG[(int) (J + LAG)] = LCHCO(MP, QTSIG, QXSIGSS);
            for (int I = 0; I < MP + 1; I++) {
                QXSIGS[J + 1][I] = QXSIGSS[I];
            }

            Calendar time = Calendar.getInstance();
            time.setTime(TM[J]);
            Integer Mth = time.get(Calendar.MONTH) + 1; // 获得降雨开始月份
            Integer Yth = time.get(Calendar.YEAR);
//            if (Yth > 2012 && Mth > 8) {
            if (Yth > 2012 && Mth > 8 && Yth0 > 2012 && Mth0 > 8) {
                QSIG[J] = QSIG[J] * 0.9;
            }
        }

        for (int i = 1; i < QSIG.length - LAG; i++) {
            if (i < QSIG0.length) {
                for (int j = 0; j < QSIG0.length; j++) {
                    if (i + j < QSIG0.length) {
                        QSIGS[i][j] = QSIG0[j + i];
                    } else {
                        QSIGS[i][j] = QSIG[i + j];
                    }
                }
            } else {
                for (int j = 0; j < QSIG0.length; j++) {
                    if (i + j < QSIG.length) {
                        QSIGS[i][j] = QSIG[i + j];
                    } else {
                        QSIGS[i][j] = QSIG[QSIG.length - 1];
                    }
                }
            }
        }

        Map<String, Object> simulationResult = new HashMap<>();

        simulationResult.put("RS", R[0]);
        simulationResult.put("RI", R[1]);
        simulationResult.put("RG", R[2]);
        simulationResult.put("QS", QS);
        simulationResult.put("QI", QI);
        simulationResult.put("QG", QG);
        simulationResult.put("QSim", QSIG);
        simulationResult.put("QGS", QGS);
        simulationResult.put("QSS", QSS);
        simulationResult.put("QIS", QIS);
        simulationResult.put("QSIGS", QSIGS);
        simulationResult.put("QXSIGS", QXSIGS);

        return simulationResult;
    }

    private double LCHCO(double MP, double RQ, double[] QX) {
        double IM = MP + 1;
        if (IM == 1.0) {
            QX[(int) (IM - 1)] = RQ;
        } else {
            for (int J = 1; J < IM; J++) {
                double Q1 = RQ;
                double Q2 = QX[J - 1];
                double Q3 = QX[J];
                QX[J - 1] = RQ;
                RQ = C0 * Q1 + C1 * Q2 + C2 * Q3;
            }
            QX[(int) (IM - 1)] = RQ;
        }
        return RQ;
    }

    private List<String> toStringList(double[] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = String.valueOf(array0[i]);
        }
        return Arrays.asList(array);
    }

    private List<String> toStringList(double[][] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = Arrays.toString(array0[i]);
        }
        return Arrays.asList(array);
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, Object> map = LAG3model();
        result.setQSim(JacksonUtil.convert(map.get("QSim"), double[].class));
        result.setRSim(new double[][]{JacksonUtil.convert(map.get("RS"), double[].class),
                JacksonUtil.convert(map.get("RI"), double[].class),
                JacksonUtil.convert(map.get("RG"), double[].class)});

        List<String> QXSIG = toStringList(JacksonUtil.convert(map.get("QXSIGS"), double[][].class));
        List<String> QSIG = toStringList(JacksonUtil.convert(map.get("QSIGS"), double[][].class));
        List<String> QSP = toStringList(JacksonUtil.convert(map.get("QSS"), double[].class));
        List<String> QIP = toStringList(JacksonUtil.convert(map.get("QIS"), double[].class));
        List<String> QGP = toStringList(JacksonUtil.convert(map.get("QGS"), double[].class));

        Map<String, List<String>> status = new HashMap<>();
        //河段下断面各时段初始流量
        status.put("QXSIG", QXSIG);
        //各时段初始总基流
        status.put("QSIG", QSIG);
        //地表坡面汇流
        status.put("QSP", QSP);
        //坡面汇流-壤中流
        status.put("QIP", QIP);
        //地下坡面汇流
        status.put("QGP", QGP);
        result.setNewStatus(status);

        return result;
    }

    public static void main(String[] args) throws Exception {
        // 读取 JSON 文件内容
        String a = new String(Files.readAllBytes(Paths.get("src/main/resources/lag_3_data.json")),
                StandardCharsets.UTF_8);
        Map data = JSON.parseObject(a, Map.class);
        LAG_3 obj = new LAG_3(data);
        System.out.println(obj.predict());
    }

}